#!/usr/bin/env python3

import os
import shlex
import subprocess
import threading

from shell_helpers import LF
import common
from thread_pool import ThreadPool

class Main(common.BuildCliFunction):
    def __init__(self):
        super().__init__(
            description='''\
Build our compiled userland examples.
'''
        )
        self.default_cstd = 'c11'
        self.default_cxxstd = 'c++17'
        self.add_argument(
            '--has-package',
            action='append',
            default=[],
            help='''\
Indicate that a given package is present in the root filesystem, which
allows us to build examples that rely on it.
''',
        )
        self.add_argument(
            '--in-tree',
            default=False,
            help='''\
Place build output inside soure tree to conveniently run it, especially when
building with the host toolchain.
''',
        )
        self.add_argument(
            '--target-cwd',
            default=False,
            help='''\
Treat targets as relative to the current working directory.
''',
        )
        self.add_argument(
            'targets',
            default=[],
            help='''\
Build only the given userland programs or all programs in the given directories.

Default: build all examples that have their package dependencies met, e.g.:
-   userland/arch/ programs only build if the target arch matches
-   an OpenBLAS example can only be built if the target root filesystem
    has the OpenBLAS libraries and headers installed, which you must inform with --has-package
''',
            nargs='*',
        )

    def _build_one(
        self,
        in_path,
        out_path,
        ccflags,
        ccflags_after=None,
        cstd=None,
        cxxstd=None,
        extra_deps=None,
        extra_objs=None,
        link=True,
    ):
        if extra_deps is None:
            extra_deps = []
        if extra_objs is None:
            extra_objs = []
        if ccflags_after is None:
            ccflags_after = []
        ret = 0
        if self.need_rebuild([in_path] + extra_objs + extra_deps, out_path):
            ccflags = ccflags.copy()
            if not link:
                ccflags.extend(['-c', LF])
            in_ext = os.path.splitext(in_path)[1]
            do_compile = True
            if in_ext in (self.env['c_ext'], self.env['asm_ext']):
                cc = self.env['gcc']
                if cstd is None:
                    std = self.default_cstd
                else:
                    std = cstd
                ccflags.extend([
                    '-fopenmp', LF,
                ])
            elif in_ext == self.env['cxx_ext']:
                cc = self.env['gxx']
                if cxxstd is None:
                    std = self.default_cxxstd
                else:
                    std = cxxstd
            else:
                do_compile = False
            if do_compile:
                os.makedirs(os.path.dirname(out_path), exist_ok=True)
                ret = self.sh.run_cmd(
                    (
                        [
                            cc, LF,
                        ] +
                        ccflags +
                        [
                            '-std={}'.format(std), LF,
                            '-o', out_path, LF,
                            in_path, LF,
                        ] +
                        self.sh.add_newlines(extra_objs) +
                        [
                            '-lm', LF,
                            '-pthread', LF,
                        ] +
                        ccflags_after
                    ),
                    extra_paths=[self.env['ccache_dir']],
                )
        return ret

    def _get_targets(self):
        if self.env['_args_given']['targets']:
            targets = self.env['targets']
            if self.env['target_cwd']:
                cwd = os.getcwd()
                targets = [os.path.join(cwd, target) for target in targets]
            return targets
        else:
            if self.env['target_cwd']:
                return [os.getcwd()]
            else:
                return [self.env['userland_source_dir']]

    def build(self):
        build_dir = self.get_build_dir()
        has_packages = set(self.env['has_package'])
        ccflags = [
            '-I', self.env['root_dir'], LF,
            '-I', self.env['userland_source_dir'], LF,
            '-O0', LF,
            '-Wall', LF,
            '-Werror', LF,
            '-Wextra', LF,
            '-Wno-unused-function', LF,
            '-ggdb3', LF,
        ]
        if self.env['static']:
            ccflags.extend(['-static', LF])
        common_obj = os.path.join(
            build_dir,
            self.env['common_basename_noext'] + self.env['obj_ext']
        )
        self._build_one(
            in_path=self.env['common_c'],
            out_path=common_obj,
            ccflags=ccflags,
            extra_deps=[self.env['common_h']],
            link=False,
        )
        common_obj_asm = os.path.join(
            build_dir,
            'arch',
            'main' + self.env['obj_ext']
        )
        common_obj_asm_relpath = os.path.join(
            'arch',
            'main' + self.env['c_ext']
        )
        self._build_one(
            in_path=os.path.join(
                self.env['userland_source_dir'],
                common_obj_asm_relpath
            ),
            out_path=common_obj_asm,
            ccflags=ccflags,
            extra_deps=[self.env['common_h']],
            link=False,
        )
        pkgs = {
            'eigen': {
                # TODO: was failing with:
                # fatal error: Eigen/Dense: No such file or directory as of
                # 975ce0723ee3fa1fea1766e6683e2f3acb8558d6
                # http://lists.busybox.net/pipermail/buildroot/2018-June/222914.html
                'ccflags': [
                    '-I',
                    os.path.join(
                        self.env['buildroot_staging_dir'],
                        'usr',
                        'include',
                        'eigen3'
                    ),
                    LF
                ],
                # Header only.
                'ccflags_after': [],
            },
            'libdrm': {},
            'openblas': {},
        }
        rootdir_abs_len = len(self.env['userland_source_dir'])
        thread_pool = ThreadPool(
            self._build_one,
            nthreads=self.env['nproc'],
        )
        class ExitLoop(Exception): pass
        try:
            for target in self._get_targets():
                target = self.resolve_userland_source(target)
                for path, in_dirnames, in_filenames in self.sh.walk(target):
                    in_dirnames.sort()
                    in_filenames.sort()
                    path_abs = os.path.abspath(path)
                    dirpath_relative_root = path_abs[rootdir_abs_len + 1:]
                    dirpath_relative_root_components = dirpath_relative_root.split(os.sep)
                    dirpath_relative_root_components_len = len(dirpath_relative_root_components)
                    do_build_dir = True
                    in_arch = False
                    if dirpath_relative_root_components_len > 0:
                        if dirpath_relative_root_components[0] == 'arch':
                            if dirpath_relative_root_components_len > 1:
                                if dirpath_relative_root_components[1] == self.env['arch']:
                                    in_arch = True
                                else:
                                    do_build_dir = False
                            else:
                                do_build_dir = False
                    if do_build_dir:
                        out_dir = os.path.join(
                            build_dir,
                            dirpath_relative_root
                        )
                        common_objs_dir = [common_obj]
                        ccflags_dir = ccflags.copy()
                        if dirpath_relative_root_components == ['gcc']:
                            cstd = 'gnu11'
                            cxxstd = 'gnu++17'
                        else:
                            cstd = self.default_cstd
                            cxxstd = self.default_cxxstd
                            # -pedantic complains even if we use -std=gnu11.
                            ccflags_dir.extend(['-pedantic', LF])
                            if in_arch:
                                ccflags_dir.extend([
                                    '-I', os.path.join(self.env['userland_source_arch_arch_dir']), LF,
                                    '-I', os.path.join(self.env['userland_source_arch_dir']), LF,
                                    '-fno-pie', LF,
                                    '-no-pie', LF,
                                ])
                                if 'freestanding' in dirpath_relative_root_components:
                                    common_objs_dir = []
                                    ccflags_dir.extend([
                                        '-ffreestanding', LF,
                                        '-nostdlib', LF,
                                        '-static', LF,
                                    ])
                                else:
                                    if 'c' in dirpath_relative_root_components:
                                        common_objs_dir = []
                                    else:
                                        common_objs_dir = [common_obj_asm]
                                if self.env['arch'] == 'arm':
                                    ccflags_dir.extend([
                                        '-Xassembler', '-mcpu=cortex-a72', LF,
                                        # To prevent:
                                        # > vfp.S: Error: selected processor does not support <FPU instruction> in ARM mode
                                        # https://stackoverflow.com/questions/41131432/cross-compiling-error-selected-processor-does-not-support-fmrx-r3-fpexc-in/52875732#52875732
                                        # We aim to take the most extended mode currently available that works on QEMU.
                                        '-Xassembler', '-mfpu=crypto-neon-fp-armv8.1', LF,
                                        '-Xassembler', '-meabi=5', LF,
                                        # Treat inline assembly as arm instead of thumb
                                        # The opposite of -mthumb.
                                        '-marm', LF,
                                        # Make gcc generate .syntax unified for inline assembly.
                                        # However, it gets ignored if -marm is given, which a GCC bug that was recently fixed:
                                        # https://stackoverflow.com/questions/54078112/how-to-write-syntax-unified-ual-armv7-inline-assembly-in-gcc/54132097#54132097
                                        # So we just write divided inline assembly for now.
                                        '-masm-syntax-unified', LF,
                                    ])
                        for in_filename in in_filenames:
                            path_relative_root = os.path.join(dirpath_relative_root, in_filename)
                            if path_relative_root == common_obj_asm_relpath:
                                continue
                            in_path = os.path.join(path, in_filename)
                            in_name, in_ext = os.path.splitext(in_filename)
                            out_path = os.path.join(
                                out_dir,
                                in_name + self.env['userland_build_ext']
                            )
                            pkg_key = in_name.split('_')[0]
                            ccflags_file = ccflags_dir.copy()
                            ccflags_after = []
                            if pkg_key in pkgs:
                                if pkg_key not in has_packages:
                                    continue
                                pkg = pkgs[pkg_key]
                                if 'ccflags' in pkg:
                                    ccflags_file.extend(pkg['ccflags'])
                                else:
                                    pkg_config_output = subprocess.check_output([
                                        self.env['buildroot_pkg_config'],
                                        '--cflags',
                                        pkg_key
                                    ]).decode()
                                    ccflags_file.extend(self.sh.shlex_split(pkg_config_output))
                                if 'ccflags_after' in pkg:
                                    ccflags_file.extend(pkg['ccflags_after'])
                                else:
                                    pkg_config_output = subprocess.check_output([
                                        self.env['buildroot_pkg_config'],
                                        '--libs',
                                        pkg_key
                                    ]).decode()
                                    ccflags_after.extend(self.sh.shlex_split(pkg_config_output))
                            error = thread_pool.submit({
                                    'in_path': in_path,
                                    'out_path': out_path,
                                    'ccflags': ccflags_file,
                                    'cstd': cstd,
                                    'cxxstd': cxxstd,
                                    'extra_objs': common_objs_dir,
                                    'ccflags_after': ccflags_after,
                            })
                            if error is not None:
                                raise ExitLoop()
        except ExitLoop:
            pass
        error = thread_pool.join()
        if error is not None:
            print(error)
            return 1
        self.sh.copy_dir_if_update(
            srcdir=build_dir,
            destdir=self.env['out_rootfs_overlay_dir'],
            filter_ext=self.env['userland_build_ext'],
        )
        return 0

    def clean(self):
        if self.env['in_tree']:
            for target in self._get_targets():
                if os.path.exists(target):
                    for path, dirnames, filenames in os.walk(target):
                        filenames.sort()
                        dirnames.sort()
                        for filename in filenames:
                            if os.path.splitext(filename)[1] in self.env['userland_out_exts']:
                                self.sh.rmrf(os.path.join(path, filename))
                else:
                    raise Exception('Path does not exist: ' + target)
        else:
            self.sh.rmrf(self.get_build_dir())

    def get_build_dir(self):
        if self.env['in_tree']:
            return self.env['userland_source_dir']
        else:
            return self.env['userland_build_dir']

if __name__ == '__main__':
    Main().cli()
